import { SiriusHost } from "~/server/api/routers/host";
import { type VulnerabilitySeverityCounts } from "~/components/VulnerabilityBarGraph";

/**
 * Categorizes vulnerabilities by severity for a host
 * @param host The host containing vulnerabilities
 * @returns Object with counts by severity level
 */
export function getVulnerabilitySeverityCounts(
  host: SiriusHost
): VulnerabilitySeverityCounts {
  // Default counts
  const counts: VulnerabilitySeverityCounts = {
    critical: 0,
    high: 0,
    medium: 0,
    low: 0,
    informational: 0,
  };

  // If host has no vulnerabilities, return empty counts
  if (!host.vulnerabilities || host.vulnerabilities.length === 0) {
    return counts;
  }

  // Count vulnerabilities by severity
  host.vulnerabilities.forEach((vulnerability) => {
    const severity = vulnerability.severity?.toLowerCase();

    if (severity === "critical") {
      counts.critical++;
    } else if (severity === "high") {
      counts.high++;
    } else if (severity === "medium") {
      counts.medium++;
    } else if (severity === "low") {
      counts.low++;
    } else if (severity === "informational" || severity === "info") {
      counts.informational++;
    }
  });

  return counts;
}

/**
 * Gets the total vulnerability count for a host
 * @param host The host containing vulnerabilities
 * @returns Total number of vulnerabilities
 */
export function getTotalVulnerabilityCount(host: SiriusHost): number {
  return host.vulnerabilities?.length || 0;
}

/**
 * Gets the count of critical vulnerabilities for a host
 * @param host The host containing vulnerabilities
 * @returns Number of critical vulnerabilities
 */
export function getCriticalVulnerabilityCount(host: SiriusHost): number {
  return (
    host.vulnerabilities?.filter(
      (v) => v.severity?.toLowerCase() === "critical"
    ).length || 0
  );
}

/**
 * Calculates an overall risk score for a host based on vulnerability severities
 * @param host The host containing vulnerabilities
 * @returns A risk score from 0-100
 */
export function calculateHostRiskScore(host: SiriusHost): number {
  if (!host.vulnerabilities || host.vulnerabilities.length === 0) {
    return 0;
  }

  // Weight each severity level
  const weights = {
    critical: 10,
    high: 5,
    medium: 2,
    low: 1,
    informational: 0.2,
  };

  const counts = getVulnerabilitySeverityCounts(host);

  // Calculate weighted score
  const score =
    counts.critical * weights.critical +
    counts.high * weights.high +
    counts.medium * weights.medium +
    counts.low * weights.low +
    counts.informational * weights.informational;

  // Cap at 100 and round to nearest integer
  return Math.min(Math.round(score), 100);
}
